#include <buola/mat.h>
#include <buola/mat/noise.h>
#include <buola/mat/cholesky.h>
#include <buola/mat/cmatfile.h>
#include <buola/chart/plots/clineplot.h>
#include <buola/chart/plots/cscatterplot.h>
#include <buola/chart/cxychart.h>
#include <buola/iterator/transform.h>

using namespace buola;

static const int sDataPoints=10000;
static const double sLengthScale=0.2;
static const double sPruneValue=1e7;
static const double sNoisePower=0.1;

struct AGaussianKernel
{
    AGaussianKernel(double pLengthScale)
        :   mLengthScale(pLengthScale)
    {}
    
    double operator()(double a,double b) const
    {
        return exp(-mLengthScale*sq(a-b));
    }
    
private:
    double mLengthScale;
};

template<typename tKernel>
mat::CVec_d rvm(const mat::CVec_d &pX,const mat::CVec_d &T,tKernel &&pKernel)
{
    int N=pX.Rows();
    int M=N+1;
    mat::CMat_d K=mat::ones(N,M);
    for(int i=0;i<N;i++)
        for(int j=0;j<N;j++)
            K(i,j)=pKernel(pX[i],pX[j]);
    
    mat::CVec_d lAlpha=mat::ones(M,1);
    double lBeta=1;
    
    mat::CVec_d lMean;
    
    for(int i=0;i<100000;i++)
    {
        start_timer();
        mat::CVec_b lValid=lAlpha<sPruneValue;
        mat::CMat_d lNewK=K(nAll,lValid);
        mat::CVec_d lNewAlpha=lAlpha(lValid,nAll);
        K=lNewK;
        lAlpha=lNewAlpha;
        M=lAlpha.size();
        
        mat::CMat_d A=mat::diagm(lAlpha);
        mat::CMat_d lSigma=(A+lBeta*(K.T()*K)).I();
        lMean=lSigma*K.T()*T;
        lMean=lBeta*lMean;
        mat::CVec_d lGamma=mat::zeros(M,1);
        for(int j=1;j<M;j++)
        {
            lGamma[j]=1-lAlpha[j]*lSigma(j,j);
            lAlpha[j]=lGamma[j]/sq(lMean[j]);
        }
        mat::CVec_d lTemp=T-K*lMean;
        lBeta=(N-sum(lGamma))/(lTemp.T()*lTemp);
        end_timer();
    }
    
    msg_info() << lMean << "\n";
        
    return K*lMean;
}

template<typename tKernel>
class CRVM
{
public:
    static constexpr double cGaussianSNR=0.1;
    static constexpr double cZeroFactor=1e-12;
    static constexpr double cMinDeltaLogAlpha=1e-3;
    static constexpr double cMinDeltaLogBeta=1e-6;
    static constexpr bool cPreferDeletion=true;
    static constexpr double cBetaMaxFactor=1e6;
    static constexpr double cAlignmentZero=1e-3;
    static constexpr double cAlignmentMax=1-cAlignmentZero;
    

    CRVM(const mat::CVec_d &pX,const mat::CVec_d &pT,const tKernel &pKernel)
        :   mX(pX)
        ,   mT(pT)
        ,   mKernel(pKernel)
        ,   mBasis(mat::ones(pX.Rows(),pX.Rows()+1))
    {
        //do initialization here
        //more or less equivalent to SB2_Initialization and whereabouts
        for(int i=0;i<mX.Rows();i++)
            for(int j=0;j<mX.Rows();j++)
                mBasis(i,j)=mKernel(mX[i],mX[j]);
            
        //this is equivalent to SB2_PreProcessBasis
        mBasisScale=norm(cols(mBasis));
        mBasisScale[mBasisScale==0]=1;
        mBasis/=extend(mBasisScale);

        ///\todo this seems to be only good for gaussian, and anyway it should be possible to specify it
        double lTStdDev=stddev(pT);
        mBeta=1/sq(lTStdDev*cGaussianSNR);

        ///\todo it should be possible to specify initial basis and additional basis... for now, just choose one
        //initialize Phi from the basis with the highest projection with the data
        mBasisTargets=mBasis.T()*pT;
    
        std::size_t lMaxProj=indmax(abs(mBasisTargets));
    
        mUsed.push_back(lMaxProj);
        mPhi=mBasis(nAll,mUsed);
    
        //mu will be initialized later, apparently, let's hope so
    
        //now initialize alpha
        mat::CVec_d dd=mPhi.T()*mPhi;
        mat::CVec_d p=mBeta*dd;
        mat::CVec_d q=mBeta*(mPhi.T()*pT);
    
        mAlpha=sq(p)/(sq(q)-p);

        mBasisPhi=mBasis.T()*mPhi;
    }
    
    enum class EAction
    {
        NONE,
        REESTIMATE,
        ADD,
        DELETE,
        TERMINATE
    };

    void Run()
    {
        ComputeStatistics();
        
        //iterate
        for(int i=0;i<5000;i++)
        {
            //find the best action
            EAction lBestAction=EAction::NONE;
            size_t lBestI=-1;       //index within all basis
            double lBestDelta=0;
            size_t lBestJ=-1;       //index within used basis
            //first see if we can delete some vector
            if(mUsed.size()>1)
            {
                for(int j=0;j<mUsed.size();j++)
                {
                    int k=mUsed[j];
                    if(mFactor[k]<=cZeroFactor)
                    {
                        double lDelta=-(sq(mQOut[k])/(mSOut[k]+mAlpha[j])-log(1+mSOut[k]/mAlpha[j]))/2;
                        if(lDelta>lBestDelta)
                        {
                            lBestAction=EAction::DELETE;
                            lBestI=k;
                            lBestDelta=lDelta;
                            lBestJ=j;
                        }
                    }
                }
            }
            //now try to add some vector
            ///\todo basis alignment support or something like that
            if(lBestAction==EAction::NONE)
            {
                mat::CVec_b lGoodFactors=mFactor>cZeroFactor;
                for(size_t k : mUsed)
                    lGoodFactors[k]=false;
                for(size_t k : mAligned|fn::values())
                    lGoodFactors[k]=false;
                for(int k=0;k<mFactor.Rows();k++)
                {
                    if(lGoodFactors[k])
                    {
                        double lQuot=sq(mQIn[k])/mSIn[k];
                        double lDelta=(lQuot-1-log(lQuot))/2.0;
                        if(lDelta>lBestDelta)
                        {
                            lBestAction=EAction::ADD;
                            lBestI=k;
                            lBestDelta=lDelta;
                        }
                    }
                }
            }
            //now try to reestimate
            if(lBestAction!=EAction::DELETE)
            {
                for(int j=0;j<mUsed.size();j++)
                {
                    int k=mUsed[j];
                    if(mFactor[k]>cZeroFactor)
                    {
                        double lNewAlpha=sq(mSOut[k])/mFactor[k];
                        double lDelta=1/lNewAlpha-1/mAlpha[j];
                        
                        lDelta=(lDelta*sq(mQIn[k])/(lDelta*mSIn[k]+1)-log(1+mSIn[k]*lDelta))/2;
                        if(lDelta>lBestDelta)
                        {
                            lBestAction=EAction::REESTIMATE;
                            lBestI=k;
                            lBestDelta=lDelta;
                            lBestJ=j;
                        }
                    }
                }
            }

            if(lBestAction==EAction::REESTIMATE)
            {
                double lNewAlpha=sq(mSOut[lBestI])/mFactor[lBestI];
                if(abs(log(lNewAlpha)-log(mAlpha[lBestJ]))<cMinDeltaLogAlpha)
                    lBestAction=EAction::NONE;
            }
            
            if(lBestAction!=EAction::NONE)
            {
                double lNewAlpha=sq(mSOut[lBestI])/mFactor[lBestI];
                
                if(lBestAction==EAction::REESTIMATE)
                {
                    double lOldAlpha=mAlpha[lBestJ];
                    mAlpha[lBestJ]=lNewAlpha;
                    mat::CVec_d lSj=mSigma(nAll,lBestJ);
                    double lDeltaInv=1/(mAlpha[lBestJ]-lOldAlpha);
                    double lKappa=1/(mSigma(lBestJ,lBestJ)+lDeltaInv);
                    mat::CVec_d lTmp=lKappa*lSj;
                    mSigma-=lTmp*lSj.T();
                    mat::CVec_d lDeltaMu=-mMu[lBestJ]*lTmp;
                    mMu+=lDeltaMu;
                    mSIn+=lKappa*sq(mBetaBasisPhi*lSj);
                    mQIn-=mBetaBasisPhi*lDeltaMu;
                }
                else if(lBestAction==EAction::ADD)
                {
                    mat::CVec_d lNewPhi=mBasis(nAll,lBestI);
                    //alignment test
                    mat::CRow_d lP=lNewPhi.T()*mPhi;
                    std::vector<size_t> lAligned=indices(lP>cAlignmentMax);
                    if(!lAligned.empty())
                    {
                        for(size_t k : lAligned)
                            mAligned.emplace(mUsed[k],lBestI);
                        continue;
                    }

                    mat::CVec_d lBasisNewPhi=mBasis.T()*lNewPhi;
                    mBasisPhi=mBasisPhi >> lBasisNewPhi;
                    mat::CVec_d lBetaNewPhi=mBeta*lNewPhi;
                    mat::CVec_d lBasisBetaNewPhi=mBeta*lBasisNewPhi;
                    mat::CVec_d lTmp=((lBetaNewPhi.T()*mPhi)*mSigma).T();
                    mAlpha=vjoin(mAlpha,lNewAlpha);
                    mPhi=mPhi >> lNewPhi;
                    double lSii=1/(lNewAlpha+mSIn[lBestI]);
                    mat::CVec_d lSi=-lSii*lTmp;
                    mat::CMat_d lTau=-lSi*lTmp.T();
                    mSigma=(mSigma+lTau) >> lSi & hjoin(lSi.T(),lSii);
                    double lMui=lSii*mQIn[lBestI];
                    mat::CVec_d lDeltaMu=vjoin(-lMui*lTmp,lMui);
                    mMu=vjoin(mMu,0)+lDeltaMu;
                    mat::CVec_d lMCi=lBasisBetaNewPhi-mBetaBasisPhi*lTmp;
                    mSIn-=lSii*sq(lMCi);
                    mQIn-=lMui*lMCi;
                    mUsed.push_back(lBestI);
                }
                else //lBestAction==EAction::DELETE
                {
                    //alignment stuff
                    mAligned.erase(lBestI);
                    
                    mBasisPhi.RemoveCol(lBestJ);
                    
                    mPhi.RemoveCol(lBestJ);
                    mAlpha.RemoveRow(lBestJ);
                    double lSjj=mSigma(lBestJ,lBestJ);
                    mat::CVec_d lSj=mSigma(nAll,lBestJ);
                    mat::CVec_d lTmp=lSj/lSjj;
                    mSigma-=lTmp*lSj.T();
                    mSigma.RemoveRow(lBestJ);
                    mSigma.RemoveCol(lBestJ);
                    mat::CVec_d lDeltaMu=-mMu[lBestJ]*lTmp;
                    double lMuj=mMu[lBestJ];
                    mMu+=lDeltaMu;
                    mMu.RemoveRow(lBestJ);
                    mat::CVec_d lJPm=mBetaBasisPhi*lSj;
                    mSIn+=sq(lJPm)/lSjj;
                    mQIn+=lJPm*lMuj/lSjj;
                    mUsed.erase(mUsed.begin()+lBestJ);
                }
                
                //do the update, since it is required for the previous three cases
                mSOut=mSIn;
                mQOut=mQIn;
                
                for(int i=0;i<mUsed.size();i++)
                {
                    int j=mUsed[i];
                    mSOut[j]=mAlpha[i]*mSIn[j]/(mAlpha[i]-mSIn[j]);
                    mQOut[j]=mAlpha[i]*mQIn[j]/(mAlpha[i]-mSIn[j]);
                }
        
                mFactor=sq(mQOut)-mSOut;
                mGamma=1-mAlpha**diag(mSigma);
                mBetaBasisPhi=mBeta*mBasisPhi;
                mLogML+=lBestDelta;
            }
            
            //recompute beta
            ///\todo only do this for the first few iterations and then every few iterations, controllable
            if(i<10||(i+1)%5==0)
            {
                double lBetaZ1=mBeta;
                mat::CVec_d lY=mPhi*mMu;
                mat::CVec_d lE=mT-lY;
                mBeta=(mX.Rows()-sum(mGamma))/double(lE.T()*lE);
                mBeta=min(mBeta,cBetaMaxFactor/var(mT));
                
                double lDeltaLogBeta=log(mBeta)-log(lBetaZ1);

                if(abs(lDeltaLogBeta)>cMinDeltaLogBeta)
                {
                    ComputeStatistics();
                }
                else if(lBestAction==EAction::NONE)
                {
                    msg_info() << "terminating after " << i << " iterations\n";
                    break;
                }
            }
        }
    }
    
    void ComputeStatistics()
    {
        mat::CMat_d A=mat::diagm(mAlpha);
        mat::CMat_d lHessian=A+mBeta*(mPhi.T()*mPhi);
        mat::CMat_d lU=mat::cholesky(lHessian).U();
        mat::CMat_d lUi=lU.I();

        mSigma=lUi*lUi.T();
        
        mat::CVec_d lPhiTargets=mPhi.T()*mT;
        mMu=(mSigma*lPhiTargets*mBeta);
 
        mat::CVec_d lY=mPhi*mMu;
        mat::CVec_d lError=mT-lY;
        double lED=lError.T()*lError;
        double lDataLikelihood=(mX.Rows()*log(mBeta)-mBeta*lED)/2;

        double lLogDetHOver2=sum(log(diag(lU)));
        mLogML=lDataLikelihood-sq(mMu).T()*mAlpha/2+sum(log(mAlpha))/2-lLogDetHOver2;
        mat::CVec_d lDiagC=sum(rows(sq(lUi)));
        mGamma=1-mAlpha**lDiagC;
        
        mBetaBasisPhi=mBeta*mBasisPhi;

        mSIn=mBeta-sum(rows(sq(mBetaBasisPhi*lUi)));
        mQIn=mBeta*(mBasisTargets-mBasisPhi*mMu);
        
        mSOut=mSIn;
        mQOut=mQIn;
        
        for(int i=0;i<mUsed.size();i++)
        {
            int j=mUsed[i];
            mSOut[j]=mAlpha[i]*mSOut[j]/(mAlpha[i]-mSIn[j]);
            mQOut[j]=mAlpha[i]*mQOut[j]/(mAlpha[i]-mSIn[j]);
        }
        
        mFactor=sq(mQOut)-mSOut;
    }
    
    mat::CVec_d GetWeights()
    {
        mat::CVec_d lReturn=mat::zeros(mT.size()+1,1);
        
        for(int i=0;i<mUsed.size();i++)
        {
            int k=mUsed[i];
            lReturn[k]=mMu[i]/mBasisScale[k];
        }
        
        return lReturn;
    }

private:
    const mat::CVec_d &mX;
    const mat::CVec_d &mT;
    const tKernel &mKernel;
    mat::CMat_d mBasis;
    mat::CRow_d mBasisScale;
    double mBeta;
    mat::CVec_d mBasisTargets;
    mat::CMat_d mBasisPhi;
    mat::CMat_d mPhi;
    std::vector<size_t> mUsed;
    std::multimap<size_t,size_t> mAligned;
    mat::CVec_d mAlpha;

    //statistics
    mat::CMat_d mSigma;
    mat::CVec_d mMu;
    double mLogML;
    mat::CVec_d mGamma;
    mat::CVec_d mSIn,mSOut;
    mat::CVec_d mQIn,mQOut;
    mat::CVec_d mFactor;
    mat::CMat_d mBetaBasisPhi;
};

template<typename tKernel>
mat::CVec_d rvm2(const mat::CVec_d &pX,const mat::CVec_d &pT,const tKernel &pKernel)
{
    start_timer();
    CRVM<tKernel> lRVM(pX,pT,pKernel);
    
    lRVM.Run();
    end_timer();
    
    mat::CVec_d lW=lRVM.GetWeights();
    
    int N=pX.Rows();
    int M=N+1;
    mat::CMat_d K=mat::ones(N,M);
    for(int i=0;i<N;i++)
        for(int j=0;j<N;j++)
            K(i,j)=pKernel(pX[i],pX[j]);
    
    return K*lW;
}

int main(int pNArg,char **pArgs)
{
    buola_init(pNArg,pArgs);

    random::engine().seed(3);
    
    try
    {
//         mat::CMatFile lFile(io::uri("/home/xavi/Documents/MATLAB/XTY.mat"));
// 
//         mat::CVec_d X=lFile.GetMat("X");
//         mat::CVec_d Y=lFile.GetMat("Y");
//         mat::CVec_d T=lFile.GetMat("T");
        
        mat::CVec_d X=mat::linspace(sDataPoints,0.0,2*M_PI);
        mat::CVec_d Y=mat::sin(X);
        mat::CVec_d T=Y+mat::wgn(Y.Rows(),Y.Cols(),sNoisePower);

        mat::CVec_d Yest=rvm2(X,T,AGaussianKernel(sLengthScale));

//        mat::CVec_d Yest=rvm(X,T,AGaussianKernel(sLengthScale));
        
        chart::PXYChart lChart=new chart::CXYChart;

        chart::PLinePlot lYPlot=new chart::CLinePlot;
        lYPlot->SetSeries(new chart::CSeries(X,Y));
        lYPlot->SetColor(img::CColor(L"darkgreen"));
        chart::PLinePlot lYestPlot=new chart::CLinePlot;
        lYestPlot->SetSeries(new chart::CSeries(X,Yest));
        lYestPlot->SetColor(img::CColor(L"blue"));
        chart::PScatterPlot lTPlot=new chart::CScatterPlot;
        lTPlot->SetSeries(new chart::CSeries(X,T));
        lTPlot->SetColor(img::CColor(L"red"));
        lTPlot->SetPointShape(chart::CPlot::CROSS);
        lTPlot->SetPointDiameter(5);
        
        lChart->AddPlot(lYPlot);
        lChart->AddPlot(lYestPlot);
        lChart->AddPlot(lTPlot);
        lChart->Map(CSize_d(800,600));
        lChart->SetTitle(L"RVM");
        
        buola_main_loop();
    }
    catch(std::exception &pE)
    {
        msg_info() << pE.what() << "\n";
    }

    return buola_finish();
}